import numpy as np
import tensorflow as tf
import os
from models.Models import LSTM_Model
from data.ProcessPoems import ProcessPoems, GenerateBatch

# 定义一些参数
tf.app.flags.DEFINE_integer('batch_size', 64, 'batch_size.')  # 每一轮唐诗数量
tf.app.flags.DEFINE_float('learning_rate', 0.01, 'learning_rate.')  # 学习率
tf.app.flags.DEFINE_string('checkpoints_dir', os.path.abspath('./checkpoints/poems/'), 'checkpoints save path')   # 训练好的模型保存起来，这样就可以利用它来生成唐诗了
tf.app.flags.DEFINE_string('file_path', os.path.abspath('./data/poems.txt'), 'file name of poems')  # 训练用的源文件路径
tf.app.flags.DEFINE_integer('epochs', 50, 'train_epoch')   # 训练多少轮
tf.app.flags.DEFINE_string('model_prefix', 'poems', 'model save prefix.')

FLAGS = tf.app.flags.FLAGS
# 唐诗预测的起始词与结束词
start_token = 'B'
end_token = 'E'


def run_training():
    if not os.path.exists(os.path.dirname(FLAGS.checkpoints_dir)):
        os.mkdir(os.path.dirname(FLAGS.checkpoints_dir))
    if not os.path.exists(FLAGS.checkpoints_dir):
        os.mkdir(FLAGS.checkpoints_dir)


    # 语料预处理
    poems_vector, word_to_int, vocabularies = ProcessPoems(FLAGS.file_path)
    # batches产生
    batch_input, batch_outputs = GenerateBatch(FLAGS.batch_size, poems_vector, word_to_int)

    input_data = tf.placeholder(tf.int32, [FLAGS.batch_size, None])
    output_targets = tf.placeholder(tf.int32, [FLAGS.batch_size, None])

    end_points = LSTM_Model(model='lstm', input_data=input_data, output_data=output_targets,
                            vocab_size=len(vocabularies), run_size=128, num_layers=2,
                            batch_size=64, learning_rate=FLAGS.learning_rate)

    saver = tf.train.Saver(tf.global_variables())
    init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
    with tf.Session() as sess:
        sess.run(init_op)
        start_epoch = 0
        checkpoint = tf.train.latest_checkpoint(FLAGS.checkpoints_dir)
        if checkpoint:
            saver.restore(sess, checkpoint)
            print("[INFO] restore from the checkpoint {0}".format(checkpoint))
            start_epoch += int(checkpoint.split('-')[-1])
        print('[INFO] start training...')
        try:
            for epoch in range(start_epoch, FLAGS.epochs):
                n = 0
                n_chunk = len(poems_vector) // FLAGS.batch_size
                for batch in range(n_chunk):
                    loss, _, _ = sess.run([
                        end_points['total_loss'],
                        end_points['last_state'],
                        end_points['train_op']
                    ], feed_dict={input_data: batch_input[n], output_targets: batch_outputs[n]})
                    n += 1
                    print('[INFO] Epoch: %d,batch:%d,training loss:%.6f' % (epoch, batch, loss))
                if epoch % 6 == 0:
                    saver.save(sess, './SaveModel/', global_step=epoch)
        except KeyboardInterrupt:
            print('[INFO] Interrupt manually, try saving checkpoint for now...')
            saver.save(sess, os.path.join(FLAGS.checkpoints_dir, FLAGS.model_prefix), global_step=epoch)
            print('[INFO] Last epoch were saved, next time will start from epoch {}.'.format(epoch))


def to_word(predict, vocabs):
    t = np.cumsum(predict)
    s = np.sum(predict)
    sample = int(np.searchsorted(t, np.random.rand(1) * s))
    if sample > len(vocabs):
        sample = len(vocabs) - 1
    return vocabs[sample]


def gen_poem(begin_word):
    batch_size = 1
    print('[INFO] loading corpus from %s' % FLAGS.file_path)
    poems_vector, word_int_map, vocabularies = ProcessPoems(FLAGS.file_path)
    input_data = tf.placeholder(tf.int32, [batch_size, None])
    end_points = LSTM_Model(model='lstm', input_data=input_data, output_data=None, vocab_size=len(vocabularies), run_size=128, num_layers=2, batch_size=64, learning_rate=FLAGS.learning_rate)
    saver = tf.train.Saver(tf.global_variables())
    init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())

    with tf.Session() as sess:
        sess.run(init_op)

        checkpoint = tf.train.latest_checkpoint('./SaveModel/')
        # saver.restore(sess, './inference/SaveModel/-24')
        #  有问题
        saver.restore(sess, './SaveModel/-30')
        x = np.array([list(map(word_int_map.get, start_token))])
        [predict, last_state] = sess.run([end_points['prediction'], end_points['last_state']],
                                         feed_dict={input_data: x})
        if begin_word:
            word = begin_word
        else:
            word = to_word(predict, vocabularies)
        poem = ''
        while word != end_token:
            print('running')
            poem += word
            x = np.zeros((1, 1))
            x[0, 0] = word_int_map[word]
            [predict, last_state] = sess.run([end_points['prediction'], end_points['last_state']],
                                             feed_dict={input_data: x, end_points['initial_state']: last_state})
            word = to_word(predict, vocabularies)
        return poem
        print('end')


def pretty_print_poem(poem):
    poem_sentences = poem.split('。')
    for s in poem_sentences:
        if s != '' and len(s) > 10:
            print(s + '。')


def main(is_train):
    if is_train:
        print("------Training------")
        run_training()
    else:
        print(".......Generating......")
        begin_word = input('请输入起始字：')      # 起始词以用户输入为准
        # begin_word = '花'      # 起始词以用户输入为准
        print("起始字TangPoems：", begin_word)
        poem2 = gen_poem(begin_word)
        # print("起始字TangPoems111111：", poem2)
        pretty_print_poem(poem2)


if __name__ == '__main__':
    tf.app.run()
